package com.shuai.robot.sse;

import cn.hutool.core.collection.CollectionUtil;
import cn.hutool.core.text.CharSequenceUtil;
import cn.hutool.json.JSONUtil;
import com.shuai.common.constant.ConversationRoleEnum;
import com.shuai.common.constant.ErrorCode;
import com.shuai.common.constant.XunFeiRobotErrorCodeEnum;
import com.shuai.common.dto.SessionChatRecordDTO;
import com.shuai.common.exception.BusinessException;
import com.shuai.robot.dto.ChatModelRps;
import com.shuai.robot.service.AIInteractService;
import com.shuai.robot.utils.ConfigContext;
import com.shuai.robot.utils.SpringBeanContext;
import lombok.SneakyThrows;
import lombok.extern.slf4j.Slf4j;
import okhttp3.Response;
import okhttp3.ResponseBody;
import okhttp3.sse.EventSource;
import okhttp3.sse.EventSourceListener;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;
import springfox.documentation.service.ApiKey;

import java.io.IOException;
import java.util.Arrays;
import java.util.Collections;
import java.util.Objects;

/**
 * @Author: huangpenglong
 * @Date: 2023/3/28 10:43
 */

@Slf4j
public class OpenAISessionChatSSEListener extends EventSourceListener {

    private static final String DONE_SIGNAL = "[DONE]";
    private final SseEmitter sseEmitter;
    private StringBuilder respContent;
    private AIInteractService aiInteractService;

    private SessionChatRecordDTO userAskRecord;

    private SessionChatRecordDTO robotRpsRecord;

    private ChatModelRps lastRps;

    private String userId;

    /**
     * 增加sessionType属性用于PDF对话释放连接
     */
    public OpenAISessionChatSSEListener(SseEmitter sseEmitter, String question, long sessionId, String userId) {
        this.userId = userId;
        userAskRecord = new SessionChatRecordDTO(sessionId, ConversationRoleEnum.USER.getRole(), question, 0);
        robotRpsRecord = new SessionChatRecordDTO(sessionId, ConversationRoleEnum.ASSISTANT.getRole(), "", 0);
        this.sseEmitter = sseEmitter;
        this.respContent = new StringBuilder();
        aiInteractService = SpringBeanContext.getBean(AIInteractService.class);
    }

    @SneakyThrows
    @Override
    public void onOpen(EventSource eventSource, Response response) {
        // 这里分来保存了用户的提问和AI的响应记录。为了查询会话中消息时能够按照特定的顺序返回
        if (userAskRecord.getSessionId() == 0) {
            // 新建会话
            long sessionId = aiInteractService.addSessionDefault(userAskRecord.getContent(), userId);
            // 更新当前会话为新建的会话
            userAskRecord.setSessionId(sessionId);
            robotRpsRecord.setSessionId(sessionId);
            // 告诉前端新建了会话
            sseEmitter.send(SseEmitter.event()
                    .id("new_session")
                    .data(sessionId)
                    .reconnectTime(3000));
        }
        aiInteractService.saveSessionBatch(Collections.singletonList(userAskRecord));
        log.info("OpenAI建立sse连接...");
    }

    @SneakyThrows
    @Override
    public void onEvent(EventSource eventSource, String id, String type, String data) {
        if (data.equals(DONE_SIGNAL)) {
            log.info("OpenAI返回数据结束,返回数据如下：\n{}\n本论问答共消耗{}token", respContent, lastRps.getUsage().getTotal_tokens());
            sseEmitter.send(SseEmitter.event()
                    .id(DONE_SIGNAL)
                    .data(DONE_SIGNAL)
                    .reconnectTime(3000));

            return;
        }

        ChatModelRps resp = JSONUtil.toBean(data, ChatModelRps.class);
        lastRps = resp;
        if (Objects.isNull(resp) || CollectionUtil.isEmpty(resp.getChoices())) {
            log.info("OpenAI返回数据异常");
            throw new BusinessException(ErrorCode.OPERATION_ERROR, "OpenAI返回数据异常");
        }

        String content = resp.getChoices().get(0).getDelta().getContent();

        if (CharSequenceUtil.isBlank(content)) {
            return;
        }
        log.info("返回数据：" + content);

        // 记录流输出结果，用于后续持久化
        this.respContent.append(content);
        // 对content数据进行处理
        content = content.replace(" ", "「」");
        content = content.replace("\n", "「·」");
        content = content.replace("\t", "「~」");
        sseEmitter.send(SseEmitter.event()
                .data(content)
                .reconnectTime(3000));
    }

    @Override
    public void onClosed(EventSource eventSource) {
        log.info("OpenAI关闭sse连接...");
        // 统计本论问答总共消耗的token数
        int totalTokens = lastRps.getUsage().getCompletion_tokens();
        robotRpsRecord.setContent(respContent.toString());
        robotRpsRecord.setTokenNum(totalTokens);
        // 扣减用户token
        aiInteractService.decreaseTokenCount(totalTokens, userId);
        // 保存会话记录
        aiInteractService.saveSessionBatch(Collections.singletonList(robotRpsRecord));
    }

    @SneakyThrows
    @Override
    public void onFailure(EventSource eventSource, Throwable t, Response response) {
        if (t != null) {
            if (t instanceof IOException) {
                log.error("网络错误！{}", t.getMessage());
            } else {
                log.error("服务错误！{}", t.getMessage());
            }
            eventSource.cancel();
            return;
        }

        if (Objects.isNull(response)) {
            log.error("OpenAI  sse连接异常:{}", "未知异常");
            eventSource.cancel();
            return;
        }

        ResponseBody body = response.body();
        if (Objects.nonNull(body)) {
            if (response.code() == 200) {
                eventSource.cancel();
                return;
            }
            String errorMsg = XunFeiRobotErrorCodeEnum.getMessageByCode(response.code());

            log.error("OpenAI  sse连接异常data：{}，异常：{}", body.string(), errorMsg);
            sseEmitter.send(SseEmitter.event()
                    .data(errorMsg)
                    .reconnectTime(3000));
        } else {
            log.error("OpenAI  sse连接异常data：{}，异常：{}", response, t);
        }
        sseEmitter.send(SseEmitter.event()
                .id(DONE_SIGNAL)
                .data(DONE_SIGNAL)
                .reconnectTime(3000));
        response.close();
        eventSource.cancel();
    }
}
